// 文件路径: dbms/execution/execution/AggregationExecutor.java
package com.dbms.execution.execution;

import com.dbms.common.Column;
import com.dbms.common.Schema;
import com.dbms.common.Tuple;
import com.dbms.common.Type;
import com.dbms.execution.plan.PlanNode;
import com.dbms.parser.ast.ExpressionNode;
import com.dbms.parser.ast.FunctionCallNode;
import com.dbms.parser.ast.GroupByNode;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 【已重构】
 * AggregationExecutor (聚合执行器) 负责处理 SQL 中的 GROUP BY 和聚合函数 (如 COUNT)。
 * 这个版本经过重构，新增了对 HAVING 子句的支持，用于在数据聚合之后进行条件过滤。
 */
public class AggregationExecutor implements Executor {

    /**
     * 子执行器，提供需要进行聚合的原始数据。
     */
    private final Executor child;

    /**
     * GROUP BY 子句的抽象语法树节点，包含了分组依据的列。
     */
    private final GroupByNode groupByNode;

    /**
     * 新增：HAVING 子句的表达式节点，用于在聚合后进行过滤。如果查询没有 HAVING 子句，则为 null。
     */
    private final ExpressionNode havingClause;

    /**
     * SELECT 查询中选择的列或聚合函数。
     */
    private final List<Object> selectColumns;

    /**
     * 聚合操作后输出的元组模式 (Schema)。
     */
    private Schema outputSchema;

    /**
     * 存储最终聚合结果的迭代器，用于 `next()` 方法逐条返回。
     */
    private Iterator<Tuple> resultIterator;

    /**
     * 构造函数。
     * @param child 子执行器
     * @param groupByNode GROUP BY 节点
     * @param selectColumns SELECT 列表
     * @param havingClause 新增：HAVING 子句节点
     */
    public AggregationExecutor(Executor child, GroupByNode groupByNode, List<Object> selectColumns, ExpressionNode havingClause) {
        this.child = child;
        this.groupByNode = groupByNode;
        this.havingClause = havingClause; // 新增
        this.selectColumns = selectColumns;
        this.resultIterator = Collections.emptyIterator(); // 初始化为空迭代器
    }

    /**
     * 打开执行器并执行所有聚合和过滤逻辑。
     * 这是执行器的核心，它会一次性从子执行器拉取所有数据，完成分组、聚合和 HAVING 过滤。
     */
    @Override
    public void open() {
        // 1. 打开子执行器并准备输出模式
        child.open();
        buildOutputSchema();

        // 2. 从子执行器获取所有元组
        List<Tuple> childTuples = new ArrayList<>();
        Tuple tuple;
        while ((tuple = child.next()) != null) {
            childTuples.add(tuple);
        }
        child.close(); // 获取完所有数据后关闭子执行器

        // 如果没有输入元组，直接返回
        if (childTuples.isEmpty()) {
            return;
        }

        // 3. 根据 GROUP BY 列对元组进行分组
        Schema inputSchema = child.getSchema();
        List<Integer> groupByIndices = groupByNode.getColumns().stream()
                .map(colName -> findColumnIndex(inputSchema, colName))
                .collect(Collectors.toList());

        Map<List<Object>, List<Tuple>> groupedTuples = childTuples.stream()
                .collect(Collectors.groupingBy(t ->
                        groupByIndices.stream()
                                .map(t::getValue)
                                .collect(Collectors.toList())
                ));

        // 4. 对每个分组进行聚合计算
        List<Tuple> aggregatedResults = new ArrayList<>();
        for (Map.Entry<List<Object>, List<Tuple>> entry : groupedTuples.entrySet()) {
            List<Object> groupKey = entry.getKey(); // 分组的键 (例如, ["Math"], ["CS"])
            List<Tuple> groupRows = entry.getValue(); // 属于该分组的所有行

            List<Object> resultValues = new ArrayList<>();
            int groupKeyIndex = 0;
            for (Object selectItem : selectColumns) {
                // 如果是 COUNT 函数
                if (selectItem instanceof FunctionCallNode && ((FunctionCallNode) selectItem).getFunctionName().equalsIgnoreCase("COUNT")) {
                    resultValues.add(groupRows.size()); // COUNT(*) 的结果就是分组中的行数
                } else if (selectItem instanceof String) { // 如果是分组的列
                    resultValues.add(groupKey.get(groupKeyIndex++)); // 从 groupKey 中获取值
                }
            }
            // 创建聚合后的元组
            aggregatedResults.add(new Tuple(resultValues, outputSchema));
        }

        // --- 【核心修改】 ---
        // 5. 在聚合之后，应用 HAVING 子句进行过滤
        if (havingClause != null) {
            // 如果存在 HAVING 子句，则过滤 aggregatedResults
            this.resultIterator = aggregatedResults.stream()
                    .filter(this::evaluateHavingClause) // 对每个聚合元组应用 having 条件
                    .iterator();
        } else {
            // 如果没有 HAVING 子句，则直接使用所有聚合结果
            this.resultIterator = aggregatedResults.iterator();
        }
        // --- 【修改结束】 ---
    }

    // --- 【新增方法】 ---
    /**
     * 对一个聚合后的元组评估 HAVING 条件。
     * 例如，对于 `HAVING COUNT(*) > 5`，这个方法会检查元组中 `COUNT(*)` 的值是否大于 5。
     * @param tuple 聚合后的元组 (其 Schema 是 outputSchema)
     * @return 如果满足 HAVING 条件则返回 true，否则返回 false
     */
    private boolean evaluateHavingClause(Tuple tuple) {
        // 1. 获取 HAVING 条件左边的值（通常是聚合函数的结果）
        String leftOperandStr = havingClause.getLeftOperand(); // 例如, "COUNT(*)"
        int leftValueIndex = -1;
        // 在输出模式中查找这个聚合函数的列索引
        for (int i = 0; i < outputSchema.getColumnCount(); i++) {
            if (outputSchema.getColumn(i).getColumnName().equalsIgnoreCase(leftOperandStr)) {
                leftValueIndex = i;
                break;
            }
        }
        if (leftValueIndex == -1) {
            throw new RuntimeException("HAVING clause aggregate function '" + leftOperandStr + "' not found in SELECT list.");
        }
        // 从元组中获取该值
        Comparable<Object> leftValue = (Comparable<Object>) tuple.getValue(leftValueIndex);

        // 2. 获取 HAVING 条件右边的常量值
        Object rightValue = havingClause.getRightOperand();

        // 3. 根据操作符执行比较
        switch (havingClause.getOperator()) {
            case "=": return leftValue.equals(rightValue);
            case ">": return leftValue.compareTo(rightValue) > 0;
            case "<": return leftValue.compareTo(rightValue) < 0;
            case ">=": return leftValue.compareTo(rightValue) >= 0;
            case "<=": return leftValue.compareTo(rightValue) <= 0;
            case "!=":
            case "<>": return !leftValue.equals(rightValue);
            default: throw new UnsupportedOperationException("Unsupported operator in HAVING clause: " + havingClause.getOperator());
        }
    }
    // --- 【新增结束】 ---

    /**
     * 根据 SELECT 列表构建输出的模式 (Schema)。
     */
    private void buildOutputSchema() {
        List<Column> columns = new ArrayList<>();
        Schema inputSchema = child.getSchema();

        for (Object item : selectColumns) {
            if (item instanceof String) {
                // 如果是普通列，从输入模式中找到对应的列信息
                columns.add(findColumn(inputSchema, (String) item));
            } else if (item instanceof FunctionCallNode) {
                // 如果是函数调用，创建一个新的列
                FunctionCallNode funcNode = (FunctionCallNode) item;
                if (funcNode.getFunctionName().equalsIgnoreCase("COUNT")) {
                    // COUNT(*) 的结果是整数类型
                    columns.add(new Column(funcNode.toString(), Type.INTEGER));
                }
            }
        }
        this.outputSchema = new Schema(columns);
    }

    /**
     * 在给定的 schema 中查找列名并返回其索引。
     * @param schema 表的模式
     * @param columnName 要查找的列名
     * @return 列的索引
     */
    private int findColumnIndex(Schema schema, String columnName) {
        for (int i = 0; i < schema.getColumnCount(); i++) {
            if (schema.getColumn(i).getColumnName().equalsIgnoreCase(columnName)) {
                return i;
            }
        }
        throw new RuntimeException("GROUP BY column '" + columnName + "' not found.");
    }

    /**
     * 在给定的 schema 中查找列名并返回其 Column 对象。
     * @param schema 表的模式
     * @param columnName 要查找的列名
     * @return Column 对象
     */
    private Column findColumn(Schema schema, String columnName) {
        for (int i = 0; i < schema.getColumnCount(); i++) {
            if (schema.getColumn(i).getColumnName().equalsIgnoreCase(columnName)) {
                return schema.getColumn(i);
            }
        }
        throw new RuntimeException("SELECT column '" + columnName + "' not found.");
    }

    /**
     * 从结果迭代器中返回下一个聚合后的元组。
     * @return 下一个元组，如果没有更多元组则返回 null。
     */
    @Override
    public Tuple next() {
        if (resultIterator.hasNext()) {
            return resultIterator.next();
        }
        return null;
    }

    /**
     * 关闭执行器。
     */
    @Override
    public void close() {}

    /**
     * 获取输出模式。
     * 如果输出模式尚未构建，则先进行构建。
     * @return 聚合后的输出模式
     */
    @Override
    public Schema getSchema() {
        if (this.outputSchema == null) {
            buildOutputSchema();
        }
        return outputSchema;
    }

    /**
     * 获取关联的计划节点。
     * @return 总是返回 null。
     */
    @Override
    public PlanNode getPlanNode() {
        return null;
    }
}